import os
from keras.models import Sequential
from keras.layers import Dense, Dropout, Flatten
from keras.layers import Conv2D, MaxPooling2D
from keras.optimizers import SGD
from keras.callbacks import ModelCheckpoint
import global_config
from keras.preprocessing.image import ImageDataGenerator
from net_model import ResNet34


def get_data_generator():
    # 声明一个ImageDataGenerator类对象，并给出你需要进行的数据增强选项
    train_datagen = ImageDataGenerator(
        # rotation_range=40,  # 角度值，0~180，影象旋轉
        # width_shift_range=0.2,  # 水平平移，相對總寬度的比例
        # height_shift_range=0.2,  # 垂直平移，相對總高度的比例
        rescale=1. / 255,
        # shear_range=0.2, # 隨機錯切換角度
        # zoom_range=0.2, # 随机缩放范围
        fill_mode='nearest',  # 填充新建立畫素的方法
        horizontal_flip=False)
    # 调用.flow_from_director()方法，第一个为数据集路径。生成数据集及标签
    train_generator = train_datagen.flow_from_directory(
        global_config.train_path,  # 目标目录
        target_size=global_config.my_config.img_size,  # 所有图片大小
        batch_size=global_config.my_config.batch_size,
        classes=list(global_config.my_config.char_set),
        color_mode='grayscale',  # grayscale”, “rbg” 之一。默认：“rgb”。图像是否被转换成 1 或 3 个颜色通道。
        class_mode='categorical')
    test_generator = train_datagen.flow_from_directory(
        global_config.test_path,  # 目标目录
        target_size=global_config.my_config.img_size,  # 所有图片大小
        batch_size=global_config.my_config.batch_size,
        classes=list(global_config.my_config.char_set),
        color_mode='grayscale',  # grayscale”, “rbg” 之一。默认：“rgb”。图像是否被转换成 1 或 3 个颜色通道。
        class_mode='categorical')
    return train_generator, test_generator


def def_model():
    '''
    模型定义
    :return:
    '''
    width, height = global_config.my_config.img_size
    model = Sequential()
    # 一层卷积层，包含了32个卷积核，大小为3*3
    model.add(Conv2D(16, (3, 3), activation='relu', input_shape=(width, height, 1)))
    model.add(Conv2D(16, (3, 3), activation='relu'))
    # 一个最大池化层，池化大小为2*2
    model.add(MaxPooling2D(pool_size=(2, 2)))
    # 遗忘层，遗忘速率为0.25
    model.add(Dropout(0.25))
    # 添加一个卷积层，包含64个卷积和，每个卷积和仍为3*3
    model.add(Conv2D(64, (3, 3), activation='relu'))
    model.add(Conv2D(64, (3, 3), activation='relu'))
    # 来一个池化层
    model.add(MaxPooling2D(pool_size=(2, 2)))
    model.add(Dropout(0.25))
    # 压平层
    model.add(Flatten())
    # 来一个全连接层
    model.add(Dense(256, activation='relu'))
    # 来一个遗忘层
    model.add(Dropout(0.5))
    # 最后为分类层
    model.add(Dense(global_config.my_config.one_captcha * len(global_config.my_config.char_set), activation='softmax'))

    sgd = SGD(lr=global_config.my_config.learning_rate, decay=1e-6, momentum=0.9, nesterov=True)
    model.compile(loss='categorical_crossentropy', optimizer=sgd, metrics=['accuracy'])
    # model.summary()
    return model


def train_model(model=None):
    '''
    训练模型
    :return:
    '''
    # 获取验证集测试集
    train_generator, test_generator = get_data_generator()

    if not model:
        # 定义模型
        model = def_model()
    # 检查点定义
    checkpoint = ModelCheckpoint(global_config.checkpoint_filepath + '/' + global_config.my_config.model_name,
                                 monitor='val_loss', save_weights_only=True, verbose=1, save_best_only=True, period=1)
    if global_config.my_config.is_load_checkpoint:
        callbacks = [checkpoint]
    else:
        callbacks = None

    if global_config.my_config.is_load_checkpoint:
        # 如果目录存在，则加载检查点模型
        if os.path.exists(global_config.checkpoint_filepath + '/' + global_config.my_config.model_name):
            model.load_weights(global_config.checkpoint_filepath + '/' + global_config.my_config.model_name)
            # 若成功加载前面保存的参数，输出下列信息
            print("模型加载")
        else:
            if not os.path.exists(global_config.checkpoint_filepath):
                os.makedirs(global_config.checkpoint_filepath)

    # 使用批量生成器擬合模型
    history = model.fit_generator(
        train_generator,
        steps_per_epoch=train_generator.n // train_generator.batch_size,
        validation_steps=test_generator.n // test_generator.batch_size,
        epochs=global_config.my_config.train_epochs,
        callbacks=callbacks,
        validation_data=test_generator)

    # 保存模型
    if global_config.my_config.is_save_model:
        if global_config.my_config.is_save_checkpoint:
            save_best_model()
        else:
            model.save(global_config.model_save_path)


def save_best_model():
    # 定义模型
    model = def_model()
    model.load_weights(global_config.checkpoint_filepath + '/' + global_config.my_config.model_name)
    model.save(global_config.model_save_path)


def get_ResNet34():
    resnet34 = ResNet34(input_shape=(global_config.my_config.img_size[0], global_config.my_config.img_size[1], 1),
                        n_classes=global_config.my_config.one_captcha * len(global_config.my_config.char_set))
    sgd = SGD(lr=0.01, decay=1e-6, momentum=0.9, nesterov=True)
    resnet34.compile(loss='categorical_crossentropy', optimizer=sgd, metrics=['accuracy'])
    return resnet34


if __name__ == '__main__':
    # model = get_ResNet34()
    model = def_model()
    train_model(model)
    # save_best_model()
